#include <stdio.h>
#include <stdlib.h>
#include "memory_allocation.h"

int64_t smallmem_lookup_table[LOOKUP_TABLE_SIZE];
int64_t smallmemcls_lookup_table[LOOKUP_TABLE_SIZE];

struct MemoryLinkedList allocated_mem;
struct MemoryLinkedList cached_mem[21];

void pm_malloc_init()
{
    allocated_mem.first = NULL;
    allocated_mem.length = 0;
    int64_t sizes[21] = {
        8,
        16,
        24,
        32,
        40,
        48,
        56,
        64,
        80,
        128,
        144,
        192,
        256,
        312,
        448,
        512,
        576,
        640,
        768,
        896,
        1024};
    int i = 0, j = 0;
    for (i = 0; i < 21; i++)
    {
        while (j <= sizes[i])
        {
            smallmem_lookup_table[j] = sizes[i];
            smallmemcls_lookup_table[j] = i;
            j += 1;
        }
    }
    for (j = 0; j < 21; j++)
    {
        cached_mem[j].first = NULL;
    }
}

// Get the memory size that should be allocated.
int64_t get_memory_to_alloc_size(int64_t orig_size)
{
    if (orig_size <= 1024)
    {
        return smallmem_lookup_table[orig_size];
    }
    else
    {
        return orig_size;
    }
}

void append_to_allocated(struct MemoryBlock *block)
{
    if (block == NULL)
    {
        printf("BLOCK IS NONE");
        exit(1);
    }
    struct MemoryBlock *allocated = allocated_mem.first;
    allocated_mem.first = block;
    allocated_mem.length += 1;
    block->next = allocated;
}

void append_to_cache(int64_t cache_cls_id, struct MemoryBlock *block)
{
    struct MemoryBlock *cache_cls = cached_mem[cache_cls_id].first;
    cached_mem->length += 1;
    cached_mem[cache_cls_id].first = block;
    cached_mem[cache_cls_id].first->next = cache_cls;
}

void *pm_malloc(int64_t size)
{
    size = get_memory_to_alloc_size(size);
    struct MemoryBlock *block;

    if (size <= 1024)
    {
        int64_t size_cls_id = smallmemcls_lookup_table[size];

        // 如果缓存中存在，则直接从缓存中取得
        if (cached_mem[size_cls_id].first != NULL)
        {
            block = cached_mem[size_cls_id].first;
            cached_mem[size_cls_id].first = cached_mem[size_cls_id].first->next;
            append_to_allocated(block);
            return block->ptr;
        }
    }

    block = (struct MemoryBlock *)malloc(sizeof(struct MemoryBlock));
    block->next = NULL;
    block->size = size;
    block->ptr = malloc((size_t)size);
    append_to_allocated(block);
    return block->ptr;
}

void pm_mfree(struct MemoryBlock *block)
{
    if (block->size <= 1024)
    {
        int64_t mem_cls_id = smallmemcls_lookup_table[block->size];
        append_to_cache(mem_cls_id, block);
    }
    else
    {
        free(block);
    }
}

void mem_cache_init(struct MemoryCache *cache)
{
    cache->stackptr = -1;
    for (int i = 0; i < CACHE_SIZE; i++)
    {
        cache->mem[i] = 0;
    }
}

int64_t mem_cache_add(struct MemoryCache *cache, void *ptr)
{
    if (cache->stackptr < CACHE_SIZE - 1)
    {
        cache->stackptr += 1;
        cache->mem[cache->stackptr] = (int64_t)ptr;
        return 1;
    }
    return 0;
}

void *mem_cache_get(struct MemoryCache *cache)
{
    void *mem;
    if (cache->stackptr >= 0)
    {
        mem = (void *)cache->mem[cache->stackptr];
        cache->stackptr -= 1;
        return mem;
    }
    else
    {
        return NULL;
    }
}

int main()
{
    pm_malloc_init();
    if (allocated_mem.first != NULL)
    {
        printf("Allocated memory should be null!");
        exit(1);
    }
    pm_malloc(1);
    if (allocated_mem.first == NULL)
    {
        printf("Allocated memory should not be null!");
        exit(1);
    }
    if (allocated_mem.first->size != 8)
    {
        printf("allocated memory size not right!");
        exit(1);
    }
    pm_malloc(191); // size should be 192
    if (allocated_mem.first->size != 192)
    {
        printf("allocated memory size not right!");
        exit(1);
    }
    pm_malloc(1999); // size should be 1999
    if (allocated_mem.first->size != 1999)
    {
        printf("allocated memory size not right!");
        exit(1);
    }
    printf("%ld\n", allocated_mem.length);
    if (allocated_mem.length != 3)
    {
        printf("length error!");
        exit(1);
    }

    struct MemoryBlock *first_block = allocated_mem.first;
    printf("size, %ld\n", allocated_mem.first->next->size);
    if (allocated_mem.first->next->size != 192)
    {
        printf("allocated memory size not right!");
        exit(1);
    }
    
    // try to remove first node.
    allocated_mem.first = allocated_mem.first->next;
    first_block->next = NULL;
    pm_mfree(first_block);

    if (allocated_mem.first->size != 192)
    {
        printf("allocated memory size not right!");
        exit(1);
    }
    printf("size, %ld\n", allocated_mem.first->size);
}
